{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "import numpy as np\n",
    "import math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sigmoid(x): \n",
    "    return 1. / (1 + np.exp(-x))\n",
    "\n",
    "# createst uniform random array w/ values in [a,b) and shape args\n",
    "def rand_arr(a, b, *args): \n",
    "    np.random.seed(0)\n",
    "    return np.random.rand(*args) * (b - a) + a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LstmParam:\n",
    "    def __init__(self, mem_cell_ct, x_dim):\n",
    "        self.mem_cell_ct = mem_cell_ct #lstm神经元的数目\n",
    "        self.x_dim = x_dim # 输入数据的维度\n",
    "        concat_len = x_dim + mem_cell_ct \n",
    "        # weight matrices\n",
    "        self.wg = rand_arr(-0.1, 0.1, mem_cell_ct, concat_len)\n",
    "        self.wi = rand_arr(-0.1, 0.1, mem_cell_ct, concat_len) \n",
    "        self.wf = rand_arr(-0.1, 0.1, mem_cell_ct, concat_len)\n",
    "        self.wo = rand_arr(-0.1, 0.1, mem_cell_ct, concat_len)\n",
    "        # bias terms\n",
    "        self.bg = rand_arr(-0.1, 0.1, mem_cell_ct) \n",
    "        self.bi = rand_arr(-0.1, 0.1, mem_cell_ct) \n",
    "        self.bf = rand_arr(-0.1, 0.1, mem_cell_ct) \n",
    "        self.bo = rand_arr(-0.1, 0.1, mem_cell_ct) \n",
    "        # diffs (derivative of loss function w.r.t. all parameters)\n",
    "        self.wg_diff = np.zeros((mem_cell_ct, concat_len)) \n",
    "        self.wi_diff = np.zeros((mem_cell_ct, concat_len)) \n",
    "        self.wf_diff = np.zeros((mem_cell_ct, concat_len)) \n",
    "        self.wo_diff = np.zeros((mem_cell_ct, concat_len)) \n",
    "        self.bg_diff = np.zeros(mem_cell_ct) \n",
    "        self.bi_diff = np.zeros(mem_cell_ct) \n",
    "        self.bf_diff = np.zeros(mem_cell_ct) \n",
    "        self.bo_diff = np.zeros(mem_cell_ct) \n",
    "\n",
    "    def apply_diff(self, lr = 1):\n",
    "        self.wg -= lr * self.wg_diff\n",
    "        self.wi -= lr * self.wi_diff\n",
    "        self.wf -= lr * self.wf_diff\n",
    "        self.wo -= lr * self.wo_diff\n",
    "        self.bg -= lr * self.bg_diff\n",
    "        self.bi -= lr * self.bi_diff\n",
    "        self.bf -= lr * self.bf_diff\n",
    "        self.bo -= lr * self.bo_diff\n",
    "        # reset diffs to zero\n",
    "        self.wg_diff = np.zeros_like(self.wg)\n",
    "        self.wi_diff = np.zeros_like(self.wi) \n",
    "        self.wf_diff = np.zeros_like(self.wf) \n",
    "        self.wo_diff = np.zeros_like(self.wo) \n",
    "        self.bg_diff = np.zeros_like(self.bg)\n",
    "        self.bi_diff = np.zeros_like(self.bi) \n",
    "        self.bf_diff = np.zeros_like(self.bf) \n",
    "        self.bo_diff = np.zeros_like(self.bo) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LstmState:\n",
    "    def __init__(self, mem_cell_ct, x_dim):\n",
    "        self.g = np.zeros(mem_cell_ct)\n",
    "        self.i = np.zeros(mem_cell_ct)\n",
    "        self.f = np.zeros(mem_cell_ct)\n",
    "        self.o = np.zeros(mem_cell_ct)\n",
    "        self.s = np.zeros(mem_cell_ct)\n",
    "        self.h = np.zeros(mem_cell_ct)\n",
    "        self.bottom_diff_h = np.zeros_like(self.h)\n",
    "        self.bottom_diff_s = np.zeros_like(self.s)\n",
    "        self.bottom_diff_x = np.zeros(x_dim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "    \n",
    "class LstmNode:\n",
    "    def __init__(self, lstm_param, lstm_state):\n",
    "        # store reference to parameters and to activations\n",
    "        self.state = lstm_state\n",
    "        self.param = lstm_param\n",
    "        # non-recurrent input to node\n",
    "        self.x = None\n",
    "        # non-recurrent input concatenated with recurrent input\n",
    "        self.xc = None\n",
    "\n",
    "    def bottom_data_is(self, x, s_prev = None, h_prev = None):\n",
    "        # if this is the first lstm node in the network\n",
    "        if s_prev == None: s_prev = np.zeros_like(self.state.s)\n",
    "        if h_prev == None: h_prev = np.zeros_like(self.state.h)\n",
    "        # save data for use in backprop\n",
    "        self.s_prev = s_prev\n",
    "        self.h_prev = h_prev\n",
    "\n",
    "        # concatenate x(t) and h(t-1)\n",
    "        xc = np.hstack((x,  h_prev))\n",
    "        self.state.g = np.tanh(np.dot(self.param.wg, xc) + self.param.bg)\n",
    "        self.state.i = sigmoid(np.dot(self.param.wi, xc) + self.param.bi)\n",
    "        self.state.f = sigmoid(np.dot(self.param.wf, xc) + self.param.bf)\n",
    "        self.state.o = sigmoid(np.dot(self.param.wo, xc) + self.param.bo)\n",
    "        self.state.s = self.state.g * self.state.i + s_prev * self.state.f\n",
    "        self.state.h = self.state.s * self.state.o\n",
    "        self.x = x\n",
    "        self.xc = xc\n",
    "    \n",
    "    def top_diff_is(self, top_diff_h, top_diff_s):\n",
    "        # top_diff_h : dL(t)/dh(t) = dl(t)/dh(t) + dL(t+1)/dh(t)\n",
    "        # top_diff_s : dL(t+1)/ds(t)\n",
    "        # notice that top_diff_s is carried along the constant error carousel（ 常量错误木马）\n",
    "        ds = self.state.o * top_diff_h + top_diff_s\n",
    "        do = self.state.s * top_diff_h\n",
    "        di = self.state.g * ds\n",
    "        dg = self.state.i * ds\n",
    "        df = self.s_prev * ds\n",
    "\n",
    "        # diffs w.r.t. vector inside sigma / tanh function\n",
    "        di_input = (1. - self.state.i) * self.state.i * di \n",
    "        df_input = (1. - self.state.f) * self.state.f * df \n",
    "        do_input = (1. - self.state.o) * self.state.o * do \n",
    "        dg_input = (1. - self.state.g ** 2) * dg\n",
    "\n",
    "        # diffs w.r.t. inputs\n",
    "        self.param.wi_diff += np.outer(di_input, self.xc)\n",
    "        self.param.wf_diff += np.outer(df_input, self.xc)\n",
    "        self.param.wo_diff += np.outer(do_input, self.xc)\n",
    "        self.param.wg_diff += np.outer(dg_input, self.xc)\n",
    "        self.param.bi_diff += di_input\n",
    "        self.param.bf_diff += df_input       \n",
    "        self.param.bo_diff += do_input\n",
    "        self.param.bg_diff += dg_input       \n",
    "\n",
    "        # compute bottom diff\n",
    "        dxc = np.zeros_like(self.xc)\n",
    "        dxc += np.dot(self.param.wi.T, di_input)\n",
    "        dxc += np.dot(self.param.wf.T, df_input)\n",
    "        dxc += np.dot(self.param.wo.T, do_input)\n",
    "        dxc += np.dot(self.param.wg.T, dg_input)\n",
    "\n",
    "        # save bottom diffs\n",
    "        self.state.bottom_diff_s = ds * self.state.f\n",
    "        self.state.bottom_diff_x = dxc[:self.param.x_dim]\n",
    "        self.state.bottom_diff_h = dxc[self.param.x_dim:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LstmNetwork():\n",
    "    def __init__(self, lstm_param):\n",
    "        self.lstm_param = lstm_param\n",
    "        self.lstm_node_list = []\n",
    "        # input sequence\n",
    "        self.x_list = []\n",
    "\n",
    "    def y_list_is(self, y_list, loss_layer):\n",
    "        \"\"\"\n",
    "        Updates diffs by setting target sequence \n",
    "        with corresponding loss layer. \n",
    "        Will *NOT* update parameters.  To update parameters,\n",
    "        call self.lstm_param.apply_diff()\n",
    "        \"\"\" \n",
    "        assert len(y_list) == len(self.x_list)\n",
    "        idx = len(self.x_list) - 1\n",
    "        # first node only gets diffs from label ...\n",
    "        loss = loss_layer.loss(self.lstm_node_list[idx].state.h, y_list[idx])\n",
    "        diff_h = loss_layer.bottom_diff(self.lstm_node_list[idx].state.h, y_list[idx])\n",
    "        # here s is not affecting loss due to h(t+1) ??????, hence we set equal to zero\n",
    "        diff_s = np.zeros(self.lstm_param.mem_cell_ct)\n",
    "        self.lstm_node_list[idx].top_diff_is(diff_h, diff_s)\n",
    "        idx -= 1\n",
    "\n",
    "        ### ... following nodes also get diffs from next nodes, hence we add diffs to diff_h\n",
    "        ### we also propagate error along constant error carousel using diff_s\n",
    "        while idx >= 0:\n",
    "            loss += loss_layer.loss(self.lstm_node_list[idx].state.h, y_list[idx])\n",
    "            diff_h = loss_layer.bottom_diff(self.lstm_node_list[idx].state.h, y_list[idx])\n",
    "            diff_h += self.lstm_node_list[idx + 1].state.bottom_diff_h # diff_h是一个递归的\n",
    "            diff_s = self.lstm_node_list[idx + 1].state.bottom_diff_s\n",
    "            self.lstm_node_list[idx].top_diff_is(diff_h, diff_s)\n",
    "            idx -= 1 \n",
    "\n",
    "        return loss\n",
    "\n",
    "    def x_list_clear(self):\n",
    "        self.x_list = []\n",
    "\n",
    "    def x_list_add(self, x):\n",
    "        self.x_list.append(x)\n",
    "        if len(self.x_list) > len(self.lstm_node_list):\n",
    "            # need to add new lstm node, create new state mem\n",
    "            lstm_state = LstmState(self.lstm_param.mem_cell_ct, self.lstm_param.x_dim)\n",
    "            self.lstm_node_list.append(LstmNode(self.lstm_param, lstm_state))\n",
    "\n",
    "        # get index of most recent x input\n",
    "        idx = len(self.x_list) - 1\n",
    "        if idx == 0:\n",
    "            # no recurrent inputs yet\n",
    "            self.lstm_node_list[idx].bottom_data_is(x)\n",
    "        else:\n",
    "            s_prev = self.lstm_node_list[idx - 1].state.s\n",
    "            h_prev = self.lstm_node_list[idx - 1].state.h\n",
    "            self.lstm_node_list[idx].bottom_data_is(x, s_prev, h_prev)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ToyLossLayer:\n",
    "    \"\"\"\n",
    "    Computes square loss with first element of hidden layer array.\n",
    "    \"\"\"\n",
    "    @classmethod\n",
    "    def loss(self, pred, label):\n",
    "        return (pred[0] - label) ** 2\n",
    "\n",
    "    @classmethod\n",
    "    def bottom_diff(self, pred, label):\n",
    "        diff = np.zeros_like(pred)\n",
    "        diff[0] = 2 * (pred[0] - label)\n",
    "        return diff"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "def example_0():\n",
    "    # learns to repeat simple sequence from random inputs\n",
    "    np.random.seed(0)\n",
    "\n",
    "    # parameters for input data dimension and lstm cell count\n",
    "    mem_cell_ct = 100\n",
    "    x_dim = 50\n",
    "    lstm_param = LstmParam(mem_cell_ct, x_dim)\n",
    "    lstm_net = LstmNetwork(lstm_param)\n",
    "    y_list = [-0.5, 0.2, 0.1, -0.5]\n",
    "    input_val_arr = [np.random.random(x_dim) for _ in y_list]\n",
    "\n",
    "    for cur_iter in range(100):\n",
    "        print(\"iter\", \"%2s\" % str(cur_iter), end=\": \")\n",
    "        for ind in range(len(y_list)):\n",
    "            lstm_net.x_list_add(input_val_arr[ind])\n",
    "\n",
    "            print(\"y_pred = [\" +\n",
    "              \", \".join([\"% 2.5f\" % lstm_net.lstm_node_list[ind].state.h[0] for ind in range(len(y_list))]) +\n",
    "              \"]\", end=\", \")\n",
    "\n",
    "        loss = lstm_net.y_list_is(y_list, ToyLossLayer)\n",
    "        print(\"loss:\", \"%.3e\" % loss)\n",
    "        lstm_param.apply_diff(lr=0.1)\n",
    "        lstm_net.x_list_clear()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "__init__() missing 2 required positional arguments: 'mem_cell_ct' and 'x_dim'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-49-09056afc90af>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0m__name__\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;34m\"__main__\"\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m     \u001b[0mexample_0\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m<ipython-input-47-7dcc78443c8c>\u001b[0m in \u001b[0;36mexample_0\u001b[0;34m()\u001b[0m\n\u001b[1;32m      7\u001b[0m     \u001b[0mx_dim\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;36m50\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      8\u001b[0m     \u001b[0mlstm_param\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mLstmParam\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmem_cell_ct\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mx_dim\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 9\u001b[0;31m     \u001b[0mlstm_net\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mLstmNetwork\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mlstm_param\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     10\u001b[0m     \u001b[0my_list\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m0.5\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m0.2\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m0.1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m-\u001b[0m\u001b[0;36m0.5\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     11\u001b[0m     \u001b[0minput_val_arr\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrandom\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrandom\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx_dim\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0m_\u001b[0m \u001b[0;32min\u001b[0m \u001b[0my_list\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-42-6e4a95fe2662>\u001b[0m in \u001b[0;36m__init__\u001b[0;34m(self, lstm_param)\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;32mclass\u001b[0m \u001b[0mLstmNetwork\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      2\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m__init__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlstm_param\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 3\u001b[0;31m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlstm_param\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mLstmParam\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      4\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlstm_node_list\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m         \u001b[0;31m# input sequence\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mTypeError\u001b[0m: __init__() missing 2 required positional arguments: 'mem_cell_ct' and 'x_dim'"
     ]
    }
   ],
   "source": [
    "if __name__ == \"__main__\":\n",
    "    example_0()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
